#! /usr/bin/python
"""
Navigation Geometry Transform Node
Transforms odometry data from Point-LIO to mavros vision pose format
"""

import rospy
from nav_msgs.msg import Odometry
from geometry_msgs.msg import PoseStamped, TransformStamped
import tf2_ros
import tf_conversions
import tf2_geometry_msgs
from math import radians

class NavGeometryTransform:
    def __init__(self):
        """Initialize the navigation geometry transform node"""
        # 获取参数
        self.input_topic = rospy.get_param('~input_topic', '/aft_mapped_to_init')
        self.output_topic = rospy.get_param('~output_topic', '/mavros/vision_pose/pose')
        self.frame_id = rospy.get_param('~frame_id', 'aircraft_base')
        self.child_frame_id = rospy.get_param('~child_frame_id', 'radar')
        
        # 变换参数
        self.translation_x = rospy.get_param('~translation_x', 0.0)
        self.translation_y = rospy.get_param('~translation_y', 0.0)
        self.translation_z = rospy.get_param('~translation_z', 0.07)  # 7cm
        self.rotation_roll = rospy.get_param('~rotation_roll', 0.0)
        self.rotation_pitch = rospy.get_param('~rotation_pitch', 0.0)
        self.rotation_yaw = rospy.get_param('~rotation_yaw', 0.0)
        
        # 发布频率控制
        self.publish_rate = rospy.get_param('~publish_rate', 50.0)  # Hz
        self.last_publish_time = rospy.Time(0)
        self.min_publish_interval = rospy.Duration(1.0 / self.publish_rate)

        rospy.loginfo("NavGeometryTransform node initialized")
        rospy.loginfo(f"Input topic: {self.input_topic}")
        rospy.loginfo(f"Output topic: {self.output_topic}")
        rospy.loginfo(f"Transform: x={self.translation_x}, y={self.translation_y}, z={self.translation_z}")
        rospy.loginfo(f"Rotation: roll={self.rotation_roll}, pitch={self.rotation_pitch}, yaw={self.rotation_yaw}")

        # 订阅输入主题
        self.odom_sub = rospy.Subscriber(self.input_topic, Odometry, self.odom_callback, queue_size=1)

        # 发布输出主题
        self.pose_pub = rospy.Publisher(self.output_topic, PoseStamped, queue_size=10)

        # 创建 tf2 缓冲和监听器
        self.tf_buffer = tf2_ros.Buffer()
        self.tf_listener = tf2_ros.TransformListener(self.tf_buffer)

        # 设置变换矩阵
        self.setup_transform()

    def setup_transform(self):
        """Setup the transformation matrix"""
        self.transform_stamped = TransformStamped()
        self.transform_stamped.header.frame_id = self.frame_id
        self.transform_stamped.child_frame_id = self.child_frame_id
        
        # 设置平移
        self.transform_stamped.transform.translation.x = self.translation_x
        self.transform_stamped.transform.translation.y = self.translation_y
        self.transform_stamped.transform.translation.z = self.translation_z

        # 创建四元数表示旋转
        quaternion = tf_conversions.transformations.quaternion_from_euler(
            radians(self.rotation_roll), 
            radians(self.rotation_pitch), 
            radians(self.rotation_yaw)
        )
        self.transform_stamped.transform.rotation.x = quaternion[0]
        self.transform_stamped.transform.rotation.y = quaternion[1]
        self.transform_stamped.transform.rotation.z = quaternion[2]
        self.transform_stamped.transform.rotation.w = quaternion[3]

    def odom_callback(self, msg):
        """Handle incoming odometry messages"""
        # 频率控制
        current_time = rospy.Time.now()
        if (current_time - self.last_publish_time) < self.min_publish_interval:
            return
        
        try:
            # 创建输入 pose
            input_pose = PoseStamped()
            input_pose.header = msg.header
            input_pose.pose = msg.pose.pose
            
            # 更新变换时间戳
            self.transform_stamped.header.stamp = current_time

            # 转换 pose
            transformed_pose = tf2_geometry_msgs.do_transform_pose(input_pose, self.transform_stamped)
            
            # 更新输出消息的时间戳和帧ID
            transformed_pose.header.stamp = current_time
            transformed_pose.header.frame_id = self.frame_id
            
            # 发布转换后的 pose
            self.pose_pub.publish(transformed_pose)
            self.last_publish_time = current_time

        except (tf2_ros.LookupException, tf2_ros.ConnectivityException, tf2_ros.ExtrapolationException) as e:
            rospy.logwarn(f"Failed to transform pose: {str(e)}")
        except Exception as e:
            rospy.logerr(f"Unexpected error in odom_callback: {str(e)}")

    def shutdown(self):
        """Clean shutdown"""
        rospy.loginfo("NavGeometryTransform node shutting down")


def main():
    """Main function"""
    try:
        # 初始化节点
        rospy.init_node('nav_geometry_transform', anonymous=False)
        
        # 创建变换节点
        transform_node = NavGeometryTransform()
        
        # 注册关闭回调
        rospy.on_shutdown(transform_node.shutdown)
        
        rospy.loginfo("NavGeometryTransform node started successfully")
        
        # 保持节点运行
        rospy.spin()
        
    except rospy.ROSInterruptException:
        rospy.loginfo("NavGeometryTransform node interrupted")
    except Exception as e:
        rospy.logerr(f"Failed to start NavGeometryTransform node: {str(e)}")


if __name__ == '__main__':
    main()
